#include <map>
#include <set>
#include <string>
#include <vector>
#include <cstdlib>
#include <iomanip>
#include <fstream>
#include <iostream>
#include <assert.h>
#include <stdlib.h>
using namespace std;
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
#include <Types.h>
#include <Defines.h>
#include <Profiler.h>
#include <DagLoader.h>
#include <MallocAlign.h>
#include <BasicDagLibrary.h>
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
union US32 {
   U32 u;
   S32 s;
};
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
Profiler prof;
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template <typename T> void PrintSpa( const map< U64, T > & spa, const string & name ) {
   typename map< U64, T >::const_iterator it;
   for( it = spa.begin(); it != spa.end(); it++ ) {
      cout << name << "[ " << it->first << " ] = " << it->second << endl;
   }
}
void PrintVec( const vector< S32 > & vec, const string & name ) {
   vector< S32 >::const_iterator it;
   U64 pos = 0ULL;
   for( it = vec.begin(); it != vec.end(); it++ ) {
      cout << name << "[ " << pos << " ] = " << *it << endl;
      pos++;
   }
}
template <typename T> void PrintSpaAsVec( const map< U32, T > & spa, const string & name ) {
   
   T prev;
   cout << name << " = [";
   typename map< U32, T >::const_iterator it;
   
   it = spa.begin();
   prev = it->second;
   it++;
   
   for( ; it != spa.end(); it++ ) {
      cout << prev << ", ";
      prev = it->second;
   }
   cout << prev << "]" << endl;

}
template <typename T> void PrintSpaAsVec( const map< S32, T > & spa, const string & name ) {
   
   T prev;
   cout << name << " = [";
   typename map< S32, T >::const_iterator it;
   
   it = spa.begin();
   prev = it->second;
   it++;
   
   for( ; it != spa.end(); it++ ) {
      cout << prev << ", ";
      prev = it->second;
   }
   cout << prev << "]" << endl;

}
void PrintSpaToMtx( const map< U64, U64 > & spa, const string & fn ) {

   LSU v;
   
   map< U64, U64 >::const_iterator it;
   it = spa.end();
   it--;
   
   const U32 m = it->first;
   const U32 n = 1;
   const U32 nnz = spa.size();
   
   ofstream file( fn.c_str() );
   
   file << "%%MatrixMarket matrix coordinate real general" << endl;
   file << m << " " << n << " " << nnz << endl;
   
   //for( U32 idx = 0; idx < nnz; idx++ ) {
   for( it = spa.begin(); it != spa.end(); it++ ) {
      v.l = it->second;
      file << (it->first + 1) << " 1 " << v.l << endl;
   }
   file.close();
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void PrintSparseArrayToMtx( const map< U64, double > & spa, const string & fmt, const MatrixProperties & mtxMetaData, const string & fn ) {

   map< U64, double >::const_iterator it;

   const U32 nnz = spa.size();
   const U32 m = RoundUpToPow2( mtxMetaData.m );
   const U32 n = RoundUpToPow2( mtxMetaData.n );
   U32 r;
   U32 c;
   U64 p;
   double v;
   
   ofstream file( fn.c_str() );
   
   file << "%%MatrixMarket matrix coordinate real general" << endl;
   file << mtxMetaData.m << " " << mtxMetaData.n << " " << nnz << endl;
   
   for( it = spa.begin(); it != spa.end(); it++ ) {
      
      p = it->first;
      v = it->second;
      
      if     ( fmt == "PSA" ) { r = Pos2RowRMA( p, m, n ); c = Pos2ColRMA( p, m, n); }
      else if( fmt == "CSR" ) { r = Pos2RowRMA( p, m, n ); c = Pos2ColRMA( p, m, n); }
      else assert( false );

      file << r << " " << c << " " << v << endl;
   }
   file.close();
   
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void CheckNumPartsVersusTreeDegree( const U32 deg, const U32 numParts ) {
   U64 logceil = LogXCeil( numParts, deg );
   U64 test = 1;
   while( logceil ) {
      test *= deg;
      logceil--;
   }
   assert( test == numParts );
}
U64 CalculateResetInterval( const U64 nnz, const U32 deg, const U32 numParts ) {

   assert( numParts > 0 );
   if( numParts == 1 ) return 0ULL;
   
   CheckNumPartsVersusTreeDegree( deg, numParts );
   
   // hgt : tree height in this tree degree
   // sub : amount to subtract from that height
   // rsi : the reset interval
   U64 effNNZ = 2 * nnz;
   U64 sub = LogXCeil( numParts, deg );
   U64 hgt = LogXCeil( effNNZ,   deg );
   U64 count = hgt - sub;
   U64 rsi = 1ULL;
   while( count ) {
      rsi *= deg;
      count--;
   }
   rsi >>= 1;
   return rsi;

}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void DeltaDecode( vector< S32 > & vec, const S32 assumedInc, const U32 resetInterval ) {

   U32 numEls = vec.size();

   vec[0] += assumedInc;
   
   for( U32 idx = 1; idx < numEls; idx++ ) {
      if( resetInterval ) {
         if( idx % resetInterval == 0 ) {
            vec[ idx ] += assumedInc;
            continue;
         }
      }
      vec[ idx ] += assumedInc + vec[ idx - 1 ];
   }
}
void Deinterleave( const U32 nnz, const map< U64, U64 > & idxspa, vector< S32 > & rowIdxs, vector< S32 > & colIdxs ) {
   
   U64 offset = 0ULL;
   U64 value  = 0ULL;
   U64 count  = 0ULL;
   const U64 mask = 0xFFFFFFFFULL;
   map< U64, U64 >::const_iterator it;
   for( it = idxspa.begin(); it != idxspa.end(); it++ ) {
      
      offset = it->first;
      value  = it->second;
      
      assert( value );
      
      while( offset > count ) {
         rowIdxs.push_back( 0 );
         colIdxs.push_back( 0 );
         count++;
      }

      US32 rowIdx;
      US32 colIdx;
      
      colIdx.u = ( value >>  0 ) & mask;
      rowIdx.u = ( value >> 32 ) & mask;

      rowIdxs.push_back( rowIdx.s );
      colIdxs.push_back( colIdx.s );
      
      count++;
   }
   while( nnz > count ) {
      rowIdxs.push_back( 0 );
      colIdxs.push_back( 0 );
      count++;
   }
   assert( count == nnz );
}
void DeinterleaveZeros( map< U64, U64 > & spa ) {
   
   map< U64, U64 > tmp = spa;
   spa.clear();
   
   map< U64, U64 >::iterator it;
   
   for( it = tmp.begin(); it != tmp.end(); it++ ) {
      U64 idx = it->first;
      U64 val = it->second;
      
      spa[ idx/2 ] = val;
   }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void IdxSpaToVector( const U32 nnz, const map< U64, U64 > & idxspa, vector< S32 > & colIdxs ) {
   
   //cout << "nnz = " << nnz << endl;
   //cout << "..." << endl;
   
   U64 mask = 0xFFFFFFFFULL;
   map< U64, U64 >::const_iterator it;
   
   U32 count = 0;
   for( it = idxspa.begin(); it != idxspa.end(); it++ ) {
      
      while( count < ( it->first << 1 ) ) {
         colIdxs.push_back( 0 );
         count++;
      }
      
      US32 v0;
      US32 v1;
      v0.u = ( it->second >> 32 ) & mask;
      v1.u = ( it->second >>  0 ) & mask;
      colIdxs.push_back( v0.s );
      colIdxs.push_back( v1.s );
      
      count++;
      count++;
   }
   
   while( count < nnz ) {
      colIdxs.push_back( 0 );
      count++;   
   }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void RMASequence( const vector< map< U64, U64 > > & spas, const U32 numParts, const MatrixProperties & mtxMetaData, const string & mtxFileName ) {

   assert( numParts == 1 );
   const map< U64, U64 > & spa = spas[0];
   map< U64, U64 >::const_iterator it;
   
   LSU v;
   U64 pos;
   U32 row;
   U32 col;
   Coordinate c;
   map< Coordinate, double > mtx;
   
   const U64 m = RoundUpToPow2( mtxMetaData.m );
   const U64 n = RoundUpToPow2( mtxMetaData.n );
   
   prof.Start( "RMA sparse array to mapped mtx" );
   for( it = spa.begin(); it != spa.end(); it++ ) {
      
      pos = it->first;
      v.l = it->second;
      
      row = Pos2RowRMA( pos, m, n );
      col = Pos2ColRMA( pos, m, n );

      c.row = row;
      c.col = col;
      mtx[ c ] = v.d;
   }
   
   prof.Stop();
   prof.Start( "PrintMappedMatrixToFile" );
   PrintMappedMatrixToFile( mtx, mtxMetaData, mtxFileName );
   prof.Stop();
}
void QTSSequence( const vector< map< U64, U64 > > & spas, const U32 numParts, const MatrixProperties & mtxMetaData, const string & mtxFileName ) {

   assert( numParts == 1 );
   const map< U64, U64 > & tmp = spas[0];
   map< U64, U64 >::const_iterator it;
   map< U64, U64 > spa = tmp;
   
   LSU v;
   U64 pos;
   U32 row;
   U32 col;
   Coordinate c;
   map< Coordinate, double > mtx;
   
   const U64 m = RoundUpToPow2( mtxMetaData.m );
   const U64 n = RoundUpToPow2( mtxMetaData.n );
   
   prof.Start( "DeinterleaveZeros" );
   DeinterleaveZeros( spa );
   prof.Stop();
   
   prof.Start( "QTS sparse array to mapped mtx" );
   for( it = spa.begin(); it != spa.end(); it++ ) {
      
      pos = it->first;
      v.l = it->second;
      
      row = Pos2RowQTS( pos, m, n );
      col = Pos2ColQTS( pos, m, n );

      c.row = row;
      c.col = col;
      mtx[ c ] = v.d;
   }
   
   prof.Stop();
   prof.Start( "PrintMappedMatrixToFile" );
   PrintMappedMatrixToFile( mtx, mtxMetaData, mtxFileName );
   prof.Stop();
}
void NZDSequence( const vector< map< U64, U64 > > & spas, const U32 numParts, const MatrixProperties & mtxMetaData, const string & mtxFileName ) {

   assert( numParts == 1 );
   assert( spas.size() == 2);
   const map< U64, U64 > & qtsSpa = spas[1];
   const map< U64, U64 > & nnzSpa = spas[0];
   map< U64, U64 >::const_iterator it;
   
   U32 row;
   U32 col;
   Coordinate c;
   map< Coordinate, double > mtx;
   
   const U64 m = RoundUpToPow2( mtxMetaData.m );
   const U64 n = RoundUpToPow2( mtxMetaData.n );
   
   U64 nzvIdx = 0ULL;
   
   prof.Start( "NZD sparse array to mapped mtx" );
   for( it = qtsSpa.begin(); it != qtsSpa.end(); it++ ) {
     
      U64 bitIdx  = 0ULL;
      U64 basePos = it->first << 6ULL;
      U64 posBits = it->second;
      while( posBits ) {
         if( posBits & 1ULL ) {
            LSU v;
            U64 pos = basePos + bitIdx;
            
            row = Pos2RowQTS( pos, m, n );
            col = Pos2ColQTS( pos, m, n );

            v.l = nnzSpa.find( nzvIdx )->second;
            nzvIdx++;
            
            c.row = row;
            c.col = col;
            mtx[ c ] = v.d;
         }
         posBits >>= 1ULL;
         bitIdx++;
      }
   }
   prof.Stop();
   prof.Start( "PrintMappedMatrixToFile" );
   PrintMappedMatrixToFile( mtx, mtxMetaData, mtxFileName );
   prof.Stop();
}
void CSRSequence( const vector< map< U64, U64 > > & spas, const U32 numParts, const MatrixProperties & mtxMetaData, const string & mtxFileName ) {
   
   assert( numParts == 1 );
   
   U64 resetInterval = CalculateResetInterval( mtxMetaData.nnz, mtxMetaData.deg, numParts );
   const map< U64, U64 > & nzvspa = spas[0];
   const map< U64, U64 > & idxspa = spas[1];
   const map< U64, U64 > & ptrspa = spas[2];
   map< U64, U64 >::const_iterator it;
   
   vector< S32 > colIdxs;
   vector< S32 > rowPtrs;
   
   S32 row = -1;
   LSU v;
   Coordinate c;
   map< Coordinate, double > mtx;
   
   prof.Start( "Column indices sparse array to vector" );
   IdxSpaToVector( mtxMetaData.nnz, idxspa, colIdxs );
   prof.Stop();
   
   prof.Start( "Row pointers sparse array to vector" );
   IdxSpaToVector( mtxMetaData.m+1, ptrspa, rowPtrs );
   prof.Stop();
   
   prof.Start( "Delta Decode" );
   DeltaDecode( colIdxs, 1, resetInterval );
   DeltaDecode( rowPtrs, 1, resetInterval );
   prof.Stop();
         
   prof.Start( "CSR vectors to mapped mtx" );
   for( it = nzvspa.begin(); it != nzvspa.end(); it++ ) {
      
      while( (U32) rowPtrs[ row+1 ] == it->first ) {
         row++;
      }
      
      c.col = colIdxs[ it->first ] + 1;
      c.row = row + 1;
      
      v.l = it->second;
      if( v.l ) {
         mtx[ c ] = v.d;
      }
   }
   prof.Stop();

   prof.Start( "PrintMappedMatrixToFile" );
   PrintMappedMatrixToFile( mtx, mtxMetaData, mtxFileName );
   prof.Stop();
   
}
void COOSequence( const vector< map< U64, U64 > > & spas, const U32 numParts, const MatrixProperties & mtxMetaData, const string & mtxFileName ) {
   
   U64 resetInterval = CalculateResetInterval( mtxMetaData.nnz, mtxMetaData.deg, numParts );
   const map< U64, U64 > & nzvspa = spas[0];
   const map< U64, U64 > & idxspa = spas[1];
   map< U64, U64 >::const_iterator it;
   
   vector< S32 > rowIdxs;
   vector< S32 > colIdxs;
      
   LSU v;
   Coordinate c;
   map< Coordinate, double > mtx;
   
   prof.Start( "Deinterleave" );
   Deinterleave( mtxMetaData.nnz, idxspa, rowIdxs, colIdxs );
   prof.Stop();

   prof.Start( "Delta Decode" );
   DeltaDecode( rowIdxs, 0, resetInterval );
   DeltaDecode( colIdxs, 1, resetInterval );
   prof.Stop();
         
   prof.Start( "COO sparse array to mapped mtx" );
   for( it = nzvspa.begin(); it != nzvspa.end(); it++ ) {
      c.row = rowIdxs[ it->first ] ;
      c.col = colIdxs[ it->first ] + 1;
      v.l = it->second;
      if( v.l ) {
         mtx[ c ] = v.d;
      }
   }
   prof.Stop();

   prof.Start( "PrintMappedMatrixToFile" );
   PrintMappedMatrixToFile( mtx, mtxMetaData, mtxFileName );
   prof.Stop();
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void Dag2Spa( map< U64, U64 > & spa, U64 pos, const U64 * nodes, U32 nodeIdx, U32 lvl, const U32 deg ) {
   
   // Dag2Spa:  dag -> sparse array

   U64 sa = Log2Ceil( deg );
   const U64 psa = 1ULL;
   U64 pmin = 0x80000000ULL;
   U64 mask = 0xFFFFFFFFULL;
   const U32 numWords = deg >> 1;
   
   if( lvl > 0 ) {
      
      U32 rptr = ( nodes[ nodeIdx ] >>  0 ) & mask;
      U32 lptr = ( nodes[ nodeIdx ] >> 32 ) & mask;

      // not a leaf
      U64 posInc = 1ULL;
      U32 lvlPow = lvl;
      if( lptr < pmin ) {
         
         while( lvlPow ) { posInc *= deg; lvlPow--; }
         posInc >>= 1ULL;
         lvl--;
         
         for( U32 wordIdx = 0; wordIdx < numWords; wordIdx++ ) {
            rptr = ( nodes[ nodeIdx + wordIdx ] >>  0 ) & mask;
            lptr = ( nodes[ nodeIdx + wordIdx ] >> 32 ) & mask;
            lptr >>= psa;
            rptr >>= psa;
            if( lptr ) {
               Dag2Spa( spa, pos, nodes, lptr, lvl, deg );
            }
            pos += posInc;
            if( rptr ) {
               Dag2Spa( spa, pos, nodes, rptr, lvl, deg );
            }
            pos += posInc;
         }
      }
      else {
         // detected a path compaction here
         U32 ptr  = lptr - pmin;
         U32 path = rptr;
         ptr >>= psa;

         while( path > sa ) {
            U32 rpath = path & ( deg - 1 );
            if( rpath ) {
               
               lvlPow = lvl;
               posInc = 1ULL;
               while( lvlPow ) { posInc *= deg; lvlPow--; }
               posInc >>= 1ULL;
               pos += rpath * posInc;
            }
            lvl--;
            path >>= sa;
         }
         Dag2Spa( spa, pos, nodes, ptr, lvl, deg );
      }
   }
   else {
      // leaf node
      U64 nonzero = 0ULL;
      for( U32 leafIdx = 0; leafIdx < numWords; leafIdx++ ) {
         U64 v = nodes[ nodeIdx + leafIdx ];
         nonzero |= v;
         if( v ) {
            spa[ pos ] = v;
         }
         pos++;
      }
      //assert( nonzero );
   }
}
void ReadDagFileToSpas( const string & mtxFileName, vector< map< U64, U64 > > & spas, MatrixProperties & mtxMetaData ) {
   
   assert( spas.size() == 0 );
   const U32 deg = mtxMetaData.deg;
   
   U64 * nodes;
   
   prof.Start( "LoadDagMtx" );
   nodes = LoadDagMtx( mtxFileName, deg, mtxMetaData );
   prof.Stop();
   
   for( U32 dagIdx = 0; dagIdx < mtxMetaData.numDags; dagIdx++ ) {
      
      map< U64, U64 > spa;
      spas.push_back( spa );
      
      U64 rootIdx = deg * mtxMetaData.rootNodeIdxs[ dagIdx ] / 2;
      U32 hgt     = mtxMetaData.dagHgts[ dagIdx ];
            
      prof.Start( "Dag2Spa" );
      Dag2Spa( spas[ dagIdx ], 0ULL, nodes, rootIdx, hgt-1, deg );
      prof.Stop();
      
   }
   
   FREE( (char*)nodes );
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
int main( int argc, char ** argv ) {

   if( argc < 6 ) {
      std::cout << "usage: " << argv[0] << " [RMA|QTS|NZD|CSR|COO] <hcdFile> <mtxFile> <deg> <numParts>" << std::endl;
      exit( -1 );
   }
   
   string mtxFmt      = argv[1];
   string hcdFileName = argv[2];
   string mtxFileName = argv[3];
   
   U32 deg      = atoi( argv[4] );
   U32 numParts = atoi( argv[5] );

   assert( hcdFileName != "stdout" );
   assert( hcdFileName != "stderr" );
   
   MatrixProperties mtxMetaData;
   vector< map< U64, U64 > > spas;
   
   mtxMetaData.deg = deg;
   ReadDagFileToSpas( hcdFileName, spas, mtxMetaData );
      
   if     ( mtxFmt == "RMA" )
      RMASequence( spas, numParts, mtxMetaData, mtxFileName );
   else if( mtxFmt == "QTS" )  QTSSequence( spas, numParts, mtxMetaData, mtxFileName );
   else if( mtxFmt == "CSR" )  CSRSequence( spas, numParts, mtxMetaData, mtxFileName );
   else if( mtxFmt == "COO" )  COOSequence( spas, numParts, mtxMetaData, mtxFileName );
   else if( mtxFmt == "NZD" )  NZDSequence( spas, numParts, mtxMetaData, mtxFileName );
   else assert( false );
         
   return 0;
}
